# coding:UTF-8
# RTK航向角数据记录程序
import serial
import time
import csv
import os
from datetime import datetime
import pynmea2
from pyproj import CRS, Transformer

# --- RTK/GPS串口参数定义 ---
rtk_params = {
    'port': '/dev/ttyUSB1',  # RTK GPS串口号
    'baudrate': 115200,
    'timeout': 0.5
}

# --- 全局变量 ---
current_heading = None  # 当前航向角
rtk_timestamp = None  # RTK时间戳

# --- RTK数据处理函数 ---
def process_hdt_data(data):
    """处理HDT语句，获取真实航向角
    Args:
        data: NMEA语句字符串
    Returns:
        float: 航向角
    """
    try:
        sentence = pynmea2.parse(data)
        if isinstance(sentence, pynmea2.HDT):
            true_heading = float(sentence.heading)
            return true_heading
    except (pynmea2.nmea.ChecksumError, pynmea2.ParseError):
        # 忽略ChecksumError和ParseError异常
        pass
    except Exception as e:
        # 捕获并打印其他异常
        print(f"处理HDT数据错误: {e}")
    return None

def read_rtk_data(ser_rtk):
    """读取RTK数据，提取航向角
    Args:
        ser_rtk: RTK串口对象
    Returns:
        heading: 航向角
        timestamp: 时间戳
    """
    global current_heading, rtk_timestamp
    
    if ser_rtk.in_waiting > 0:
        data = ser_rtk.readline().decode('utf-8', errors='ignore').rstrip()
        print(f"原始数据: {data}")  # 打印原始NMEA数据，便于调试
        heading = process_hdt_data(data)
        if heading is not None:
            current_heading = heading
            rtk_timestamp = datetime.now().strftime('%H:%M:%S.%f')
            return current_heading, rtk_timestamp
    
    return current_heading, rtk_timestamp

# --- 主程序 ---
if __name__ == '__main__':
    try:
        # 打开RTK串口
        ser_rtk = serial.Serial(**rtk_params)
        print(f"RTK串口已打开: {ser_rtk.is_open}, 端口: {rtk_params['port']}")
        
        # 创建CSV文件
        filename_input = input("请输入保存文件的名称(如输入数字则保存为'rtk_数字.csv'，输入'q'退出): ")
        if filename_input.lower() == 'q':
            print("程序退出")
            exit(0)
            
        filename = f"rtk_{filename_input}.csv"
        print(f"将保存数据到文件: {filename}")
        
        with open(filename, 'w', newline='') as csvfile:
            csv_writer = csv.writer(csvfile)
            csv_writer.writerow([
                'Timestamp', 
                'RTK_Heading', 
                'RTK_Timestamp'
            ])
            
            # 设置运行时间
            run_time = 10.0  # 5秒
            start_time = time.time()
            data_count = 0
            
            print(f"开始记录数据，将运行{run_time}秒...")
            
            # 主循环
            while time.time() - start_time < run_time:
                # 读取RTK数据
                heading, rtk_time = read_rtk_data(ser_rtk)
                
                # 如果读取到有效的航向角数据
                if heading is not None and rtk_time is not None:
                    # 显示RTK数据
                    print(f"RTK航向角: {heading}°, 时间: {rtk_time}")
                    
                    # 记录数据
                    timestamp = time.time() - start_time
                    csv_writer.writerow([
                        timestamp,
                        heading,
                        rtk_time
                    ])
                    csvfile.flush()
                    data_count += 1
                    
                    if data_count % 10 == 0:
                        elapsed = time.time() - start_time
                        remaining = max(0, run_time - elapsed)
                        print(f"已记录 {data_count} 组数据, 剩余时间: {remaining:.1f}秒")
                
                # 短暂休眠以减少CPU使用率
                time.sleep(0.01)
            
            print(f"共记录 {data_count} 组数据")
            
            # 添加等待用户输入，让用户选择继续记录或退出
            while True:
                next_action = input("记录完成。输入新的文件名继续记录，或输入'q'退出: ")
                if next_action.lower() == 'q':
                    break
                    
                # 新文件名
                filename = f"rtk_{next_action}.csv"
                print(f"将保存数据到新文件: {filename}")
                
                with open(filename, 'w', newline='') as csvfile:
                    csv_writer = csv.writer(csvfile)
                    csv_writer.writerow([
                        'Timestamp', 
                        'RTK_Heading', 
                        'RTK_Timestamp'
                    ])
                    
                    # 重置计数和时间
                    start_time = time.time()
                    data_count = 0
                    
                    print(f"开始记录新数据，将运行{run_time}秒...")
                    
                    # 主循环
                    while time.time() - start_time < run_time:
                        # 读取RTK数据
                        heading, rtk_time = read_rtk_data(ser_rtk)
                        
                        # 如果读取到有效的航向角数据
                        if heading is not None and rtk_time is not None:
                            # 显示RTK数据
                            print(f"RTK航向角: {heading}°, 时间: {rtk_time}")
                            
                            # 记录数据
                            timestamp = time.time() - start_time
                            csv_writer.writerow([
                                timestamp,
                                heading,
                                rtk_time
                            ])
                            csvfile.flush()
                            data_count += 1
                            
                            if data_count % 10 == 0:
                                elapsed = time.time() - start_time
                                remaining = max(0, run_time - elapsed)
                                print(f"已记录 {data_count} 组数据, 剩余时间: {remaining:.1f}秒")
                        
                        # 短暂休眠以减少CPU使用率
                        time.sleep(0.01)
                    
                    print(f"共记录 {data_count} 组数据")
    
    except KeyboardInterrupt:
        print("\n程序被用户中断")
    
    except serial.SerialException as e:
        print(f"串口错误: {e}")
    
    except Exception as e:
        print(f"程序错误: {e}")
    
    finally:
        # 关闭串口
        if 'ser_rtk' in locals() and ser_rtk.is_open:
            ser_rtk.close()
            print("RTK串口已关闭") 